import backtrader as bt
import logging
import numpy as np
import pandas as pd
from datetime import datetime


class BaseStrategy(bt.Strategy):
    """
    基础策略类，提供通用功能和日志记录
    """
    params = (
        ('log_trades', True),  # 是否记录交易日志
        ('log_data', False),  # 是否记录数据详情
        ('min_data_required', 30),  # 最小数据要求
    )

    def __init__(self):
        """
        初始化策略
        """
        super().__init__()
        self.logger = logging.getLogger(self.__class__.__name__)
        self.logger.info(f"初始化 {self.__class__.__name__} 策略")

        # 初始化状态变量
        self.orders = {}
        self.indicators_created = False
        self.asset_names = {}

        # 记录资产名称
        for i, data in enumerate(self.datas):
            self.asset_names[data] = data._name if hasattr(data, '_name') else f"Data{i}"
            self.logger.info(f"数据源 {i}: {self.asset_names[data]}")

            # 记录数据长度
            self.logger.debug(f"{self.asset_names[data]} 初始数据长度: {len(data)}")

        # 验证数据源
        self._validate_data_sources()

    def _validate_data_sources(self):
        """
        验证数据源是否足够
        """
        self.logger.info(f"共有 {len(self.datas)} 个数据源")

        # 检查是否有足够的数据源
        if len(self.datas) == 0:
            self.logger.error("策略需要至少一个数据源")
            raise ValueError("策略需要至少一个数据源")

        # 检查数据长度是否满足最小要求
        for data in self.datas:
            asset_name = self.asset_names[data]
            if len(data) < self.params.min_data_required:
                self.logger.warning(
                    f"{asset_name} 数据不足: {len(data)} 条，需要至少 {self.params.min_data_required} 条")

    def _create_indicators(self):
        """
        创建技术指标（由子类实现）
        """
        raise NotImplementedError("子类必须实现 _create_indicators 方法")

    def next(self):
        """
        每个bar执行一次的策略逻辑（由子类实现）
        """
        raise NotImplementedError("子类必须实现 next 方法")

    def notify_order(self, order):
        """
        处理订单状态通知
        """
        if order.status in [order.Submitted, order.Accepted]:
            # 订单已提交/接受，无需处理
            return

        if order.status in [order.Completed]:
            # 订单已完成
            asset_name = self.asset_names[order.data]

            if order.isbuy():
                self.logger.info(f"买入完成: {asset_name}, 价格: {order.executed.price:.2f}, "
                                 f"数量: {order.executed.size}, 成本: {order.executed.value:.2f}, "
                                 f"佣金: {order.executed.comm:.2f}")
            elif order.issell():
                self.logger.info(f"卖出完成: {asset_name}, 价格: {order.executed.price:.2f}, "
                                 f"数量: {order.executed.size}, 收入: {order.executed.value:.2f}, "
                                 f"佣金: {order.executed.comm:.2f}")

            # 从订单字典中移除
            if order in self.orders:
                del self.orders[order]

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            # 订单问题
            asset_name = self.asset_names[order.data] if order.data else "未知资产"
            status = order.getstatusname()
            self.logger.warning(f"订单问题 ({status}): {asset_name}, 原因: {order.getstatusname()}")

            # 从订单字典中移除
            if order in self.orders:
                del self.orders[order]

    def notify_trade(self, trade):
        """
        处理交易通知
        """
        if not trade.isclosed:
            return

        asset_name = self.asset_names[trade.data]
        self.logger.info(f"交易关闭: {asset_name}, 利润: {trade.pnl:.2f}, 净利润: {trade.pnlcomm:.2f}")

    def stop(self):
        """
        回测结束时调用
        """
        self.logger.info("回测结束")
        self.logger.info(f"最终投资组合价值: {self.broker.getvalue():.2f}")

        # 打印各资产持仓情况
        self.logger.info("最终持仓:")
        for data in self.datas:
            asset_name = self.asset_names[data]
            position = self.getposition(data)
            if position.size != 0:
                value = position.size * data.close[0]
                self.logger.info(f"  {asset_name}: {position.size} 股, 价值: {value:.2f}")


class AllWeatherStrategy(BaseStrategy):
    """
    修正后的全天候投资策略

    该策略基于Ray Dalio的全天候投资组合理念，通过配置多元化的资产类别
    并在特定时间进行再平衡，实现在各种经济环境下的稳健收益。

    修复点:
    1. 添加了数据验证和错误处理
    2. 修复了目标配置资产缺失时的处理
    3. 添加了价格零值保护
    4. 优化了再平衡逻辑
    """

    params = (
        ('rebalance_months', [1, 4, 7, 10]),  # 每年的1月、4月、7月、10月再平衡
        ('target_allocation', {
            'SPY': 0.30,  # 股票
            'TLT': 0.40,  # 长期债券
            'IEF': 0.15,  # 中期债券
            'GLD': 0.075,  # 黄金
            'DBC': 0.075  # 大宗商品
        }),
        ('sma_period', 120),  # 移动平均线周期
        ('min_trade_value', 100),  # 最小交易价值
        ('price_zero_threshold', 0.01),  # 价格零值阈值
    )

    def __init__(self):
        """
        初始化策略
        """
        super().__init__()

        # 初始化指标字典
        self.sma_indicators = {}

        # 创建指标
        self._create_indicators()

    def _create_indicators(self):
        """
        创建技术指标
        """
        self.logger.info("创建技术指标...")

        # 为每个数据源创建SMA指标
        for data in self.datas:
            asset_name = self.asset_names[data]

            # 确保有足够数据计算指标
            if len(data) < self.params.sma_period:
                self.logger.warning(f"{asset_name} 数据不足 ({len(data)} < {self.params.sma_period})，跳过SMA创建")
                continue

            # 创建SMA指标
            self.sma_indicators[data] = bt.indicators.SimpleMovingAverage(
                data.close, period=self.params.sma_period
            )
            self.logger.info(f"为 {asset_name} 创建了 {self.params.sma_period} 周期SMA")

        self.indicators_created = True
        self.logger.info("指标创建完成")

    def next(self):
        """
        每个bar执行一次的策略逻辑
        """
        # 如果指标尚未创建（因为数据不足），尝试创建
        if not self.indicators_created:
            # 检查所有数据源是否都有足够数据
            min_data_required = self.params.sma_period
            all_data_ready = all(len(data) >= min_data_required for data in self.datas)

            if all_data_ready:
                self._create_indicators()
            else:
                # 如果数据不足，跳过此bar
                self.logger.debug("数据不足，跳过策略逻辑")
                return

        # 获取当前日期
        current_date = self.data.datetime.date(0)
        current_month = current_date.month

        # 检查是否在再平衡月份
        if current_month in self.params.rebalance_months:
            self.logger.info(f"再平衡月份: {current_date}")
            self._rebalance_portfolio()

    def _rebalance_portfolio(self):
        """
        执行投资组合再平衡
        """
        self.logger.info("开始投资组合再平衡")

        # 计算投资组合总价值
        portfolio_value = self.broker.getvalue()
        self.logger.info(f"投资组合总价值: {portfolio_value:.2f}")

        # 遍历所有资产
        for data in self.datas:
            asset_name = self.asset_names[data]

            # 获取目标配置比例（默认为0）
            target_percent = self.params.target_allocation.get(asset_name, 0.0)
            if target_percent <= 0:
                self.logger.debug(f"{asset_name} 无目标配置，跳过")
                continue

            # 获取当前价格，检查零值
            current_price = data.close[0]
            if current_price < self.params.price_zero_threshold:
                self.logger.warning(f"{asset_name} 当前价格 {current_price:.4f} 过低，跳过交易")
                continue

            # 计算目标价值
            target_value = portfolio_value * target_percent
            self.logger.info(f"{asset_name} 目标价值: {target_value:.2f} ({target_percent * 100:.1f}%)")

            # 获取当前持仓价值
            current_value = self.broker.getvalue([data])
            self.logger.info(f"{asset_name} 当前价值: {current_value:.2f}")

            # 计算价值差异
            value_diff = target_value - current_value
            abs_diff = abs(value_diff)

            # 如果差异小于最小交易价值，跳过
            if abs_diff < self.params.min_trade_value:
                self.logger.info(
                    f"{asset_name} 价值差异 {abs_diff:.2f} 小于最小交易价值 {self.params.min_trade_value:.2f}，跳过调整")
                continue

            # 确定交易方向
            if value_diff > 0:
                # 买入
                size = int(value_diff // current_price)
                if size > 0:
                    self.logger.info(f"买入 {asset_name}: {size} 股 @ {current_price:.2f}")
                    self.buy(data=data, size=size)
                else:
                    self.logger.warning(f"计算买入数量为0: 价值差异={value_diff:.2f}, 价格={current_price:.2f}")
            else:
                # 卖出
                size = int(abs_diff // current_price)
                if size > 0:
                    self.logger.info(f"卖出 {asset_name}: {size} 股 @ {current_price:.2f}")
                    self.sell(data=data, size=size)
                else:
                    self.logger.warning(f"计算卖出数量为0: 价值差异={value_diff:.2f}, 价格={current_price:.2f}")


class SimpleMovingAverageCrossover(BaseStrategy):
    """
    简单移动平均线交叉策略
    用于调试和基本测试
    """
    params = (
        ('fast_period', 10),  # 快速移动平均线周期
        ('slow_period', 30),  # 慢速移动平均线周期
        ('trade_size', 100),  # 每次交易数量
    )

    def __init__(self):
        """
        初始化策略
        """
        super().__init__()

        # 确保只有一个数据源
        if len(self.datas) > 1:
            self.logger.warning("该策略设计为单资产使用，将使用第一个数据源")

        self.data = self.datas[0]
        self.asset_name = self.asset_names[self.data]

        # 创建指标
        self.fast_sma = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.fast_period
        )
        self.slow_sma = bt.indicators.SimpleMovingAverage(
            self.data.close, period=self.params.slow_period
        )

        # 交叉信号
        self.crossover = bt.indicators.CrossOver(
            self.fast_sma, self.slow_sma
        )

        self.logger.info(
            f"为 {self.asset_name} 创建了 {self.params.fast_period}/{self.params.slow_period} SMA 交叉策略")

    def next(self):
        """
        每个bar执行一次的策略逻辑
        """
        # 检查是否有持仓
        position = self.getposition(self.data)

        # 金叉信号：快速线上穿慢速线
        if self.crossover > 0:
            if not position:
                self.logger.info(f"金叉信号: 买入 {self.asset_name}")
                self.buy(data=self.data, size=self.params.trade_size)

        # 死叉信号：快速线下穿慢速线
        elif self.crossover < 0:
            if position:
                self.logger.info(f"死叉信号: 卖出 {self.asset_name}")
                self.sell(data=self.data, size=position.size)


class BuyAndHold(BaseStrategy):
    """
    买入并持有策略
    用于基准测试
    """
    params = (
        ('trade_size', 100),  # 每次交易数量
    )

    def __init__(self):
        """
        初始化策略
        """
        super().__init__()

        # 确保只有一个数据源
        if len(self.datas) > 1:
            self.logger.warning("该策略设计为单资产使用，将使用第一个数据源")

        self.data = self.datas[0]
        self.asset_name = self.asset_names[self.data]
        self.order_executed = False

    def next(self):
        """
        每个bar执行一次的策略逻辑
        """
        # 只在第一天买入
        if not self.order_executed:
            self.logger.info(f"初始买入: {self.asset_name}")
            self.buy(data=self.data, size=self.params.trade_size)
            self.order_executed = True


class PortfolioRebalance(BaseStrategy):
    """
    定期再平衡策略
    简化版的全天候策略，不依赖技术指标
    """
    params = (
        ('rebalance_months', [1, 7]),  # 再平衡月份 (1月和7月)
        ('target_allocation', {
            'SPY': 0.60,
            'TLT': 0.40
        }),
        ('min_trade_value', 100),  # 最小交易价值
    )

    def next(self):
        """
        每个bar执行一次的策略逻辑
        """
        # 获取当前日期
        current_date = self.data.datetime.date(0)
        current_month = current_date.month

        # 检查是否在再平衡月份
        if current_month in self.params.rebalance_months:
            self.logger.info(f"再平衡月份: {current_date}")
            self._rebalance_portfolio()

    def _rebalance_portfolio(self):
        """
        执行投资组合再平衡
        """
        self.logger.info("开始投资组合再平衡")

        # 计算投资组合总价值
        portfolio_value = self.broker.getvalue()
        self.logger.info(f"投资组合总价值: {portfolio_value:.2f}")

        # 遍历所有资产
        for data in self.datas:
            asset_name = self.asset_names[data]

            # 获取目标配置比例（默认为0）
            target_percent = self.params.target_allocation.get(asset_name, 0.0)
            if target_percent <= 0:
                self.logger.debug(f"{asset_name} 无目标配置，跳过")
                continue

            # 获取当前价格
            current_price = data.close[0]

            # 计算目标价值
            target_value = portfolio_value * target_percent
            self.logger.info(f"{asset_name} 目标价值: {target_value:.2f} ({target_percent * 100:.1f}%)")

            # 获取当前持仓价值
            current_value = self.broker.getvalue([data])
            self.logger.info(f"{asset_name} 当前价值: {current_value:.2f}")

            # 计算价值差异
            value_diff = target_value - current_value
            abs_diff = abs(value_diff)

            # 如果差异小于最小交易价值，跳过
            if abs_diff < self.params.min_trade_value:
                self.logger.info(
                    f"{asset_name} 价值差异 {abs_diff:.2f} 小于最小交易价值 {self.params.min_trade_value:.2f}，跳过调整")
                continue

            # 确定交易方向
            if value_diff > 0:
                # 买入
                size = int(value_diff // current_price)
                if size > 0:
                    self.logger.info(f"买入 {asset_name}: {size} 股 @ {current_price:.2f}")
                    self.buy(data=data, size=size)
            else:
                # 卖出
                size = int(abs_diff // current_price)
                if size > 0:
                    self.logger.info(f"卖出 {asset_name}: {size} 股 @ {current_price:.2f}")
                    self.sell(data=data, size=size)


# 策略注册表，便于动态加载
STRATEGY_REGISTRY = {
    'allweather': AllWeatherStrategy,
    'sma_crossover': SimpleMovingAverageCrossover,
    'buyhold': BuyAndHold,
    'rebalance': PortfolioRebalance,
}


def get_strategy(strategy_name):
    """
    根据名称获取策略类

    参数:
        strategy_name: 策略名称

    返回:
        class: 策略类
    """
    strategy_class = STRATEGY_REGISTRY.get(strategy_name.lower())
    if not strategy_class:
        raise ValueError(f"未知策略: {strategy_name}. 可用策略: {', '.join(STRATEGY_REGISTRY.keys())}")
    return strategy_class